# coding: utf-8
"""
File: matching.py
Project: AlphaPose
File Created: Monday, 1st October 2018 12:53:12 pm
Author: Yuliang Xiu (yuliangxiu@sjtu.edu.cn)
Copyright 2018 - 2018 Shanghai Jiao Tong University, Machine Vision and Intelligence Group
"""

import argparse
import os

import cv2
from tqdm import tqdm


def generate_fake_cor(img, out_path):
    print("Generate fake correspondence files...%s" % out_path)
    fd = open(out_path, "w")
    height, width, channels = img.shape

    for x in range(width):
        for y in range(height):
            ret = fd.write("%d %d %d %d %f \n" % (x, y, x, y, 1.0))
    fd.close()


def orb_matching(img1_path, img2_path, vidname, img1_id, img2_id):
    out_path = "%s/%s_%s_orb.txt" % (vidname, img1_id, img2_id)
    # print(out_path)

    img1 = cv2.cvtColor(cv2.imread(img1_path), cv2.COLOR_BGR2RGB)
    img2 = cv2.cvtColor(cv2.imread(img2_path), cv2.COLOR_BGR2RGB)

    # Initiate ORB detector
    orb = cv2.ORB_create(nfeatures=10000, scoreType=cv2.ORB_FAST_SCORE)

    # find the keypoints and descriptors with ORB
    kp1, des1 = orb.detectAndCompute(img1, None)
    kp2, des2 = orb.detectAndCompute(img2, None)

    if len(kp1) * len(kp2) < 400:
        generate_fake_cor(img1, out_path)
        return

    # FLANN parameters
    FLANN_INDEX_LSH = 6
    index_params = dict(algorithm=FLANN_INDEX_LSH,
                        table_number=12,  # 12
                        key_size=12,  # 20
                        multi_probe_level=2)  # 2

    search_params = dict(checks=100)  # or pass empty dictionary

    flann = cv2.FlannBasedMatcher(index_params, search_params)

    matches = flann.knnMatch(des1, des2, k=2)

    # Open file
    fd = open(out_path, "w")

    # ratio test as per Lowe's paper
    for i, m_n in enumerate(matches):
        if len(m_n) != 2:
            continue
        elif m_n[0].distance < 0.80 * m_n[1].distance:
            ret = fd.write("%d %d %d %d %f \n" % (
                kp1[m_n[0].queryIdx].pt[0], kp1[m_n[0].queryIdx].pt[1], kp2[m_n[0].trainIdx].pt[0], kp2[m_n[0].trainIdx].pt[1], m_n[0].distance))

    # Close opened file
    fd.close()

    # print(os.stat(out_path).st_size)

    # if os.stat(out_path).st_size < 1000:
    #     generate_fake_cor(img1, out_path)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='FoseFlow Matching')
    parser.add_argument('--orb', type=int, default=0)
    args = parser.parse_args()

    image_dir = "posetrack_data/images"
    imgnames = []
    vidnames = []

    for a, b, c in os.walk(image_dir):
        if len(a.split("/")) == 4:
            vidnames.append(a)

    for vidname in tqdm(sorted(vidnames)):
        for a, b, c in os.walk(vidname):
            c = [item for item in c if "jpg" in item]
            imgnames = sorted(c)
            break
        for imgname in imgnames[:-1]:
            if 'crop' in imgname:
                continue
            img1 = os.path.join(vidname, imgname)
            len_name = len(imgname.split(".")[0])
            if len_name == 5:
                img2 = os.path.join(vidname, "%05d.jpg" % (int(imgname.split(".")[0]) + 1))
            else:
                img2 = os.path.join(vidname, "%08d.jpg" % (int(imgname.split(".")[0]) + 1))
            if not os.path.exists(img2):
                continue
            img1_id = img1.split(".")[0].split("/")[-1]
            img2_id = img2.split(".")[0].split("/")[-1]
            if args.orb:
                cor_file = "%s/%s_%s_orb.txt" % (vidname, img1_id, img2_id)
            else:
                cor_file = "%s/%s_%s.txt" % (vidname, img1_id, img2_id)
            if not os.path.exists(cor_file) or os.stat(cor_file).st_size < 1000:
                if args.orb:
                    # calc orb matching
                    orb_matching(img1, img2, vidname, img1_id, img2_id)
                else:
                    # calc deep matching
                    cmd = "./deepmatching/deepmatching %s %s -nt 10 -downscale 3 -out %s/%s_%s.txt > cache" % (img1, img2, vidname, img1_id, img2_id)
                    os.system(cmd)
